import SimpleITK as sitk 
from monai.networks.nets.segresnet import SegResNet
import torch 
from segmentation.nestedformer.nested_former import NestedFormer
# pretrain_model_path = "./models/best_model_0.7105_0.pt"

# def load_segment_model():
#     model = SegResNet(3, 16, 2, 2)
#     model.load_state_dict(torch.load(pretrain_model_path, map_location="cpu"))
#     return model

# pretrain_model_path = "./models/final_model_0.8997.pt"
pretrain_model_path = "./models/final_model_0.8410.pt"

def load_segment_model():
    model = NestedFormer(model_num=2,
                        out_channels=3,
                        image_size=(24, 192, 192),
                        window_size=(3, 6, 6),
                        pool_size=((1, 2, 2), (1, 2, 2), (2, 2, 2), (2, 2, 2)))
    model.load_state_dict(torch.load(pretrain_model_path, map_location="cpu"))
    return model